# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

load("//build_tools/bazel:build_defs.oss.bzl", "iree_cmake_extra_content")
load("//build_tools/bazel:iree_bitcode_library.bzl", "iree_bitcode_library", "iree_link_bitcode")

package(
    default_visibility = ["//visibility:public"],
    features = ["layering_check"],
    licenses = ["notice"],  # Apache 2.0
)

#===------------------------------------------------------------------------===#
# UKernel bitcode files
#===------------------------------------------------------------------------===#

iree_cmake_extra_content(
    content = """
iree_compiler_targeting_iree_arch(_IREE_UKERNEL_BITCODE_BUILD_ARM_64 "arm_64")
if(_IREE_UKERNEL_BITCODE_BUILD_ARM_64)
""",
    inline = True,
)

# All headers transitively included by code in this directory. Bazel-only.
UKERNEL_ARM_64_INTERNAL_HEADERS = [
    "common_arm_64.h",
    "common_arm_64_entry_point.h",
    "mmt4d_arm_64_internal.h",
    "pack_arm_64_internal.h",
    "unpack_arm_64_internal.h",
    "//runtime/src/iree/builtins/ukernel:internal_headers_filegroup",
    "//runtime/src/iree/schemas:cpu_data_headers_filegroup",
]

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_entry_points",
    srcs = [
        "mmt4d_arm_64_entry_point.c",
        "pack_arm_64_entry_point.c",
        "query_tile_sizes_arm_64_entry_point.c",
        "unpack_arm_64_entry_point.c",
    ],
    # wasm_64 here is a proxy for "some reasonable 64-bit architecture". This
    # should match the `ukernel_bitcode_64bit_base` bitcode library, because we
    # want this code inlined into the ukernel entry points so that code path
    # selection logic can evaluate at compile time and unused code paths can be
    # DCE'd. The entry points don't need to be architecture-specialized, all
    # they do is return function pointers to the actual architecture-specialized
    # symbols defined in the other iree_bitcode_library's below.
    arch = "wasm_64",
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_base",
    srcs = [
        "mmt4d_arm_64.c",
        "pack_arm_64.c",
        "unpack_arm_64.c",
    ],
    arch = "arm_64",
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_fullfp16",
    srcs = ["mmt4d_arm_64_fullfp16.c"],
    arch = "arm_64",
    copts = ["-march=armv8.2-a+fp16"],
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_fp16fml",
    srcs = ["mmt4d_arm_64_fp16fml.c"],
    arch = "arm_64",
    copts = ["-march=armv8.2-a+fp16fml"],
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_bf16",
    srcs = ["mmt4d_arm_64_bf16.c"],
    arch = "arm_64",
    copts = ["-march=armv8.2-a+bf16"],
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_dotprod",
    srcs = ["mmt4d_arm_64_dotprod.c"],
    arch = "arm_64",
    copts = ["-march=armv8.2-a+dotprod"],
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_bitcode_library(
    name = "ukernel_bitcode_arm_64_i8mm",
    srcs = ["mmt4d_arm_64_i8mm.c"],
    arch = "arm_64",
    copts = ["-march=armv8.2-a+i8mm"],
    internal_hdrs = UKERNEL_ARM_64_INTERNAL_HEADERS,
)

iree_link_bitcode(
    name = "ukernel_bitcode_arm_64",
    bitcode_files = [
        "ukernel_bitcode_arm_64_base.bc",
        "ukernel_bitcode_arm_64_fullfp16.bc",
        "ukernel_bitcode_arm_64_fp16fml.bc",
        "ukernel_bitcode_arm_64_bf16.bc",
        "ukernel_bitcode_arm_64_dotprod.bc",
        "ukernel_bitcode_arm_64_i8mm.bc",
    ],
)

iree_cmake_extra_content(
    content = """
elseif(IREE_BUILD_COMPILER AND IREE_TARGET_BACKEND_LLVM_CPU)
iree_make_empty_file("${CMAKE_CURRENT_BINARY_DIR}/ukernel_bitcode_arm_64.bc")
iree_make_empty_file("${CMAKE_CURRENT_BINARY_DIR}/ukernel_bitcode_arm_64_entry_points.bc")
endif()  # _IREE_UKERNEL_BITCODE_BUILD_ARM_64
""",
    inline = True,
)
